## Check MCMC convergence by examining R-hat, effective sample size, 
## Bayesian fraction of missing information, and traceplots. 
# Takes a long time + a lot of memory to run this script.


library(rstan)
library(dplyr)
library(ggpubr)
library(ggplot2)
library(ggthemes)
library(ggridges)
library(stringr)
if(!require(wpmarble)){
  devtools::install_github("wpmarble/wpmarble")
  library(wpmarble)
}

theme_set(theme_bw() + theme(panel.spacing = unit(0, "in")))


source("Code/hlm/functions/stan_utility.R")
source("Code/hlm/functions/summarize_results.R")

select <- dplyr::select
extract <- rstan::extract



## Load samples 
load("Data/demos_socioecon_pid7_full_sample.rdata")




## Some diagnostics to assess convergence -- see source code in Code/hlm/functions/**
check_div(sampout)
check_energy(sampout)
check_treedepth(sampout)


# Get summaries for each param to look at some diagnostics
summ = readRDS("Data/summary_socioecon_pid7.rds")


# look at Rhats - should be 1 (at convergence). Rule of thumb is <1.2 OK
summary(summ$Rhat)

# effective sample size (# of independent draws from the posterior)
summary(summ$n_eff)
rm(summ);gc()

# make some trace plots 
sd_educ_tp = traceplot(sampout,   "educ_sd")
sd_hieduc_tp = traceplot(sampout, "hieduc_sd")
sd_gov_tp = traceplot(sampout, "gov_sd")
sd_invest_tp = traceplot(sampout, "invest_sd")
sd_workers_tp = traceplot(sampout, "workers_sd")
sd_local_tp = traceplot(sampout, "local_sd")
sd_intercept_tp = traceplot(sampout, "intercept_sd")
ggsave(sd_educ_tp, file = "Output/extras/traceplot_sd_educ.pdf", width=11, height=8.5)
ggsave(sd_hieduc_tp, file = "Output/extras/traceplot_sd_hieduc.pdf", width=11, height=8.5)
ggsave(sd_gov_tp, file = "Output/extras/traceplot_sd_gov.pdf", width=11, height=8.5)
ggsave(sd_invest_tp, file = "Output/extras/traceplot_sd_invest.pdf", width=11, height=8.5)
ggsave(sd_workers_tp, file = "Output/extras/traceplot_sd_workers.pdf", width=11, height=8.5)
ggsave(sd_local_tp, file = "Output/extras/traceplot_sd_local.pdf", width=11, height=8.5)
ggsave(sd_intercept_tp, file = "Output/extras/traceplot_sd_intercept.pdf", width=11, height=8.5)

lp_tp = traceplot(sampout, "lp__")
ggsave(lp_tp, file = "Output/extras/traceplot_lp.pdf", width=11, height=8.5)


# remove ggplot objects
to_rm = ls()[grepl("_tp$",ls())]
rm(list=to_rm); gc()

tp_gamma_educ = traceplot(sampout, "gamma_educ", nrow = 4)
ggsave(tp_gamma_educ, file = "Output/extras/traceplot_gamma_educ.pdf", width=14, height=7)

tp_gamma_gov = traceplot(sampout, "gamma_gov", nrow = 6)
ggsave(tp_gamma_gov, file = "Output/extras/traceplot_gamma_gov.pdf", width=10, height=7)

tp_gamma_intercept = traceplot(sampout, "gamma_intercept")
ggsave(tp_gamma_intercept, file = "Output/extras/traceplot_gamma_intercept.pdf", width=10, height=7)

tp_intercept = traceplot(sampout, paste0("intercept[", sample(1:standat$n_resp, 12), "]"))
ggsave(tp_intercept, file = "Output/extras/traceplot_intercept.pdf", width=10, height=10)

tp_educ_sd = traceplot(sampout, "educ_sd")
ggsave(tp_educ_sd, file = "Output/extras/traceplot_educ_sd.pdf", width=10, height=10)

# pick some random betas to plot traceplot
tp_betas_invest = traceplot(sampout, paste0("beta_invest[", sample(1:standat$n_resp, 12), ",1]"))
ggsave(tp_betas_invest, file = "Output/extras/traceplot_beta_invest.pdf", width=10, height=10)

tp_betas_hieduc = traceplot(sampout, paste0("beta_hieduc[", sample(1:standat$n_resp, 12), ",1]"))
ggsave(tp_betas_hieduc, file = "Output/extras/traceplot_beta_hieduc.pdf", width=10, height=10)



